from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
from bson import ObjectId
from pymongo.collection import Collection
import hashlib
import json
from utils.logger import get_logger
from utils.exceptions import (
    AuditException,
    ComplianceException,
    ValidationException
)

logger = get_logger(__name__)

class RegulatoryAuditGateway:
    """监管审计网关服务 (RAG)
    
    负责:
    1. 合规性检查
    2. 审计日志记录
    3. 监管报告生成
    4. 数据隐私保护
    5. 法规遵循验证
    """
    
    def __init__(self, mongo_db):
        self.db = mongo_db
        self.audit_logs_collection: Collection = mongo_db.audit_logs
        self.compliance_rules_collection: Collection = mongo_db.compliance_rules
        self.regulatory_reports_collection: Collection = mongo_db.regulatory_reports
        self.privacy_records_collection: Collection = mongo_db.privacy_records
        
        # 合规规则配置
        self.compliance_rules = {
            'transaction_limits': {
                'max_amount': 1000000,  # 最大交易金额
                'daily_limit': 5000000,  # 日交易限额
                'monthly_limit': 50000000  # 月交易限额
            },
            'kyc_requirements': {
                'required_fields': ['name', 'id_number', 'phone', 'email'],
                'verification_level': 'enhanced'
            },
            'aml_rules': {
                'suspicious_amount_threshold': 100000,
                'velocity_check_enabled': True,
                'blacklist_check_enabled': True
            },
            'data_retention': {
                'audit_log_retention_days': 2555,  # 7年
                'transaction_data_retention_days': 1825,  # 5年
                'user_data_retention_days': 1095  # 3年
            }
        }
        
        # 创建索引
        self._create_indexes()
        
        # 初始化合规规则
        self._initialize_compliance_rules()
    
    def _create_indexes(self):
        """创建数据库索引"""
        try:
            # 审计日志索引
            self.audit_logs_collection.create_index([("timestamp", -1)])
            self.audit_logs_collection.create_index("user_id")
            self.audit_logs_collection.create_index("action_type")
            self.audit_logs_collection.create_index("chain_id")
            self.audit_logs_collection.create_index("transaction_hash")
            
            # 合规规则索引
            self.compliance_rules_collection.create_index("rule_type")
            self.compliance_rules_collection.create_index("status")
            
            # 监管报告索引
            self.regulatory_reports_collection.create_index("report_type")
            self.regulatory_reports_collection.create_index("period")
            self.regulatory_reports_collection.create_index("created_at")
            
            # 隐私记录索引
            self.privacy_records_collection.create_index("user_id")
            self.privacy_records_collection.create_index("data_type")
            self.privacy_records_collection.create_index("access_timestamp")
            
        except Exception as e:
            logger.warning(f"Failed to create indexes: {str(e)}")
    
    def _initialize_compliance_rules(self):
        """初始化合规规则"""
        try:
            # 检查是否已存在规则
            existing_rules = self.compliance_rules_collection.count_documents({})
            if existing_rules > 0:
                return
            
            # 创建默认合规规则
            default_rules = [
                {
                    'rule_id': 'txn_limit_001',
                    'rule_type': 'transaction_limit',
                    'name': '交易金额限制',
                    'description': '单笔交易金额不得超过100万',
                    'parameters': {'max_amount': 1000000},
                    'status': 'active',
                    'created_at': datetime.utcnow(),
                    'updated_at': datetime.utcnow()
                },
                {
                    'rule_id': 'kyc_001',
                    'rule_type': 'kyc_verification',
                    'name': 'KYC身份验证',
                    'description': '用户必须完成增强级身份验证',
                    'parameters': {
                        'required_fields': ['name', 'id_number', 'phone', 'email'],
                        'verification_level': 'enhanced'
                    },
                    'status': 'active',
                    'created_at': datetime.utcnow(),
                    'updated_at': datetime.utcnow()
                },
                {
                    'rule_id': 'aml_001',
                    'rule_type': 'anti_money_laundering',
                    'name': '反洗钱检查',
                    'description': '检查可疑交易和黑名单',
                    'parameters': {
                        'suspicious_amount_threshold': 100000,
                        'velocity_check_enabled': True,
                        'blacklist_check_enabled': True
                    },
                    'status': 'active',
                    'created_at': datetime.utcnow(),
                    'updated_at': datetime.utcnow()
                }
            ]
            
            self.compliance_rules_collection.insert_many(default_rules)
            logger.info("Default compliance rules initialized")
            
        except Exception as e:
            logger.error(f"Failed to initialize compliance rules: {str(e)}")
    
    def log_audit_event(self, user_id: str, action_type: str, details: Dict[str, Any], 
                       chain_id: str = None, transaction_hash: str = None) -> str:
        """记录审计事件"""
        try:
            # 生成审计ID
            audit_id = self._generate_audit_id(user_id, action_type)
            
            # 创建审计日志
            audit_log = {
                'audit_id': audit_id,
                'user_id': user_id,
                'action_type': action_type,
                'details': details,
                'chain_id': chain_id,
                'transaction_hash': transaction_hash,
                'timestamp': datetime.utcnow(),
                'ip_address': details.get('ip_address'),
                'user_agent': details.get('user_agent'),
                'session_id': details.get('session_id'),
                'result': details.get('result', 'success'),
                'error_message': details.get('error_message'),
                'data_hash': self._calculate_data_hash(details)
            }
            
            result = self.audit_logs_collection.insert_one(audit_log)
            
            logger.info(f"Audit event logged: {audit_id} - {action_type} by {user_id}")
            return audit_id
            
        except Exception as e:
            logger.error(f"Failed to log audit event: {str(e)}")
            raise AuditException(f"Failed to log audit event: {str(e)}")
    
    def _generate_audit_id(self, user_id: str, action_type: str) -> str:
        """生成审计ID"""
        timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S')
        data = f"{user_id}_{action_type}_{timestamp}"
        hash_value = hashlib.md5(data.encode()).hexdigest()[:8]
        return f"AUDIT_{timestamp}_{hash_value}"
    
    def _calculate_data_hash(self, data: Dict[str, Any]) -> str:
        """计算数据哈希值"""
        try:
            # 移除敏感字段
            filtered_data = {k: v for k, v in data.items() 
                           if k not in ['password', 'private_key', 'secret']}
            
            data_str = json.dumps(filtered_data, sort_keys=True, default=str)
            return hashlib.sha256(data_str.encode()).hexdigest()
        except Exception:
            return "hash_error"
    
    def check_compliance(self, operation_type: str, data: Dict[str, Any]) -> Dict[str, Any]:
        """检查合规性"""
        try:
            compliance_result = {
                'compliant': True,
                'violations': [],
                'warnings': [],
                'recommendations': []
            }
            
            # 获取相关合规规则
            rules = self._get_applicable_rules(operation_type)
            
            for rule in rules:
                violation = self._check_rule_compliance(rule, data)
                if violation:
                    compliance_result['violations'].append(violation)
                    compliance_result['compliant'] = False
            
            # 记录合规检查结果
            self._log_compliance_check(operation_type, data, compliance_result)
            
            return compliance_result
            
        except Exception as e:
            logger.error(f"Failed to check compliance: {str(e)}")
            raise ComplianceException(f"Failed to check compliance: {str(e)}")
    
    def _get_applicable_rules(self, operation_type: str) -> List[Dict[str, Any]]:
        """获取适用的合规规则"""
        try:
            # 根据操作类型映射规则类型
            rule_mapping = {
                'transaction': ['transaction_limit', 'anti_money_laundering'],
                'account_creation': ['kyc_verification'],
                'contract_deployment': ['code_audit', 'security_check'],
                'data_access': ['privacy_protection', 'data_retention']
            }
            
            rule_types = rule_mapping.get(operation_type, [])
            
            if not rule_types:
                return []
            
            cursor = self.compliance_rules_collection.find({
                'rule_type': {'$in': rule_types},
                'status': 'active'
            })
            
            return list(cursor)
            
        except Exception as e:
            logger.error(f"Failed to get applicable rules: {str(e)}")
            return []
    
    def _check_rule_compliance(self, rule: Dict[str, Any], data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """检查单个规则的合规性"""
        try:
            rule_type = rule['rule_type']
            parameters = rule['parameters']
            
            if rule_type == 'transaction_limit':
                return self._check_transaction_limit(parameters, data)
            elif rule_type == 'kyc_verification':
                return self._check_kyc_compliance(parameters, data)
            elif rule_type == 'anti_money_laundering':
                return self._check_aml_compliance(parameters, data)
            elif rule_type == 'privacy_protection':
                return self._check_privacy_compliance(parameters, data)
            
            return None
            
        except Exception as e:
            logger.error(f"Failed to check rule compliance: {str(e)}")
            return {
                'rule_id': rule.get('rule_id', 'unknown'),
                'violation_type': 'check_error',
                'message': f"Failed to check rule: {str(e)}"
            }
    
    def _check_transaction_limit(self, parameters: Dict[str, Any], data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """检查交易限额"""
        amount = data.get('amount', 0)
        max_amount = parameters.get('max_amount', 1000000)
        
        if amount > max_amount:
            return {
                'rule_id': 'txn_limit_001',
                'violation_type': 'transaction_limit_exceeded',
                'message': f"Transaction amount {amount} exceeds limit {max_amount}",
                'severity': 'high'
            }
        
        return None
    
    def _check_kyc_compliance(self, parameters: Dict[str, Any], data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """检查KYC合规性"""
        required_fields = parameters.get('required_fields', [])
        user_data = data.get('user_data', {})
        
        missing_fields = [field for field in required_fields if not user_data.get(field)]
        
        if missing_fields:
            return {
                'rule_id': 'kyc_001',
                'violation_type': 'kyc_incomplete',
                'message': f"Missing required KYC fields: {', '.join(missing_fields)}",
                'severity': 'medium'
            }
        
        return None
    
    def _check_aml_compliance(self, parameters: Dict[str, Any], data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """检查反洗钱合规性"""
        amount = data.get('amount', 0)
        threshold = parameters.get('suspicious_amount_threshold', 100000)
        
        if amount > threshold:
            return {
                'rule_id': 'aml_001',
                'violation_type': 'suspicious_transaction',
                'message': f"Transaction amount {amount} exceeds AML threshold {threshold}",
                'severity': 'high'
            }
        
        return None
    
    def _check_privacy_compliance(self, parameters: Dict[str, Any], data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """检查隐私保护合规性"""
        # 检查是否访问了敏感数据
        sensitive_fields = ['private_key', 'password', 'ssn', 'bank_account']
        accessed_fields = data.get('accessed_fields', [])
        
        sensitive_accessed = [field for field in accessed_fields if field in sensitive_fields]
        
        if sensitive_accessed and not data.get('authorized', False):
            return {
                'rule_id': 'privacy_001',
                'violation_type': 'unauthorized_sensitive_access',
                'message': f"Unauthorized access to sensitive fields: {', '.join(sensitive_accessed)}",
                'severity': 'critical'
            }
        
        return None
    
    def _log_compliance_check(self, operation_type: str, data: Dict[str, Any], result: Dict[str, Any]):
        """记录合规检查结果"""
        try:
            self.log_audit_event(
                user_id=data.get('user_id', 'system'),
                action_type='compliance_check',
                details={
                    'operation_type': operation_type,
                    'compliance_result': result,
                    'data_summary': self._create_data_summary(data)
                }
            )
        except Exception as e:
            logger.error(f"Failed to log compliance check: {str(e)}")
    
    def _create_data_summary(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """创建数据摘要（移除敏感信息）"""
        sensitive_fields = ['password', 'private_key', 'secret', 'token']
        
        summary = {}
        for key, value in data.items():
            if key in sensitive_fields:
                summary[key] = '[REDACTED]'
            elif isinstance(value, (dict, list)):
                summary[key] = f"[{type(value).__name__}]"
            else:
                summary[key] = str(value)[:100]  # 限制长度
        
        return summary
    
    def generate_regulatory_report(self, report_type: str, period: str, 
                                 start_date: datetime = None, end_date: datetime = None) -> str:
        """生成监管报告"""
        try:
            if not start_date:
                start_date = datetime.utcnow() - timedelta(days=30)
            if not end_date:
                end_date = datetime.utcnow()
            
            # 生成报告数据
            report_data = self._collect_report_data(report_type, start_date, end_date)
            
            # 创建报告文档
            report_doc = {
                'report_id': self._generate_report_id(report_type, period),
                'report_type': report_type,
                'period': period,
                'start_date': start_date,
                'end_date': end_date,
                'data': report_data,
                'summary': self._generate_report_summary(report_data),
                'created_at': datetime.utcnow(),
                'created_by': 'system',
                'status': 'completed'
            }
            
            result = self.regulatory_reports_collection.insert_one(report_doc)
            report_id = report_doc['report_id']
            
            logger.info(f"Regulatory report generated: {report_id}")
            return report_id
            
        except Exception as e:
            logger.error(f"Failed to generate regulatory report: {str(e)}")
            raise AuditException(f"Failed to generate regulatory report: {str(e)}")
    
    def _generate_report_id(self, report_type: str, period: str) -> str:
        """生成报告ID"""
        timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S')
        return f"RPT_{report_type.upper()}_{period.upper()}_{timestamp}"
    
    def _collect_report_data(self, report_type: str, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """收集报告数据"""
        try:
            if report_type == 'transaction_summary':
                return self._collect_transaction_summary(start_date, end_date)
            elif report_type == 'compliance_violations':
                return self._collect_compliance_violations(start_date, end_date)
            elif report_type == 'audit_activity':
                return self._collect_audit_activity(start_date, end_date)
            elif report_type == 'user_activity':
                return self._collect_user_activity(start_date, end_date)
            else:
                return {'error': f'Unknown report type: {report_type}'}
                
        except Exception as e:
            logger.error(f"Failed to collect report data: {str(e)}")
            return {'error': str(e)}
    
    def _collect_transaction_summary(self, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """收集交易摘要数据"""
        # 这里应该从实际的交易数据中统计
        # 目前返回模拟数据
        return {
            'total_transactions': 15420,
            'total_volume': 125000000,
            'average_transaction_size': 8108,
            'largest_transaction': 950000,
            'suspicious_transactions': 23,
            'failed_transactions': 156,
            'success_rate': 98.99
        }
    
    def _collect_compliance_violations(self, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """收集合规违规数据"""
        try:
            pipeline = [
                {
                    '$match': {
                        'timestamp': {'$gte': start_date, '$lte': end_date},
                        'action_type': 'compliance_check',
                        'details.compliance_result.compliant': False
                    }
                },
                {
                    '$group': {
                        '_id': '$details.compliance_result.violations.violation_type',
                        'count': {'$sum': 1}
                    }
                }
            ]
            
            violations = list(self.audit_logs_collection.aggregate(pipeline))
            
            return {
                'total_violations': sum(v['count'] for v in violations),
                'violation_types': {v['_id']: v['count'] for v in violations},
                'period_start': start_date.isoformat(),
                'period_end': end_date.isoformat()
            }
            
        except Exception as e:
            logger.error(f"Failed to collect compliance violations: {str(e)}")
            return {'error': str(e)}
    
    def _collect_audit_activity(self, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """收集审计活动数据"""
        try:
            pipeline = [
                {
                    '$match': {
                        'timestamp': {'$gte': start_date, '$lte': end_date}
                    }
                },
                {
                    '$group': {
                        '_id': '$action_type',
                        'count': {'$sum': 1}
                    }
                }
            ]
            
            activities = list(self.audit_logs_collection.aggregate(pipeline))
            total_events = self.audit_logs_collection.count_documents({
                'timestamp': {'$gte': start_date, '$lte': end_date}
            })
            
            return {
                'total_audit_events': total_events,
                'activity_breakdown': {a['_id']: a['count'] for a in activities},
                'period_start': start_date.isoformat(),
                'period_end': end_date.isoformat()
            }
            
        except Exception as e:
            logger.error(f"Failed to collect audit activity: {str(e)}")
            return {'error': str(e)}
    
    def _collect_user_activity(self, start_date: datetime, end_date: datetime) -> Dict[str, Any]:
        """收集用户活动数据"""
        try:
            pipeline = [
                {
                    '$match': {
                        'timestamp': {'$gte': start_date, '$lte': end_date}
                    }
                },
                {
                    '$group': {
                        '_id': '$user_id',
                        'activity_count': {'$sum': 1},
                        'last_activity': {'$max': '$timestamp'}
                    }
                }
            ]
            
            user_activities = list(self.audit_logs_collection.aggregate(pipeline))
            
            return {
                'active_users': len(user_activities),
                'total_user_activities': sum(u['activity_count'] for u in user_activities),
                'average_activities_per_user': sum(u['activity_count'] for u in user_activities) / len(user_activities) if user_activities else 0,
                'period_start': start_date.isoformat(),
                'period_end': end_date.isoformat()
            }
            
        except Exception as e:
            logger.error(f"Failed to collect user activity: {str(e)}")
            return {'error': str(e)}
    
    def _generate_report_summary(self, report_data: Dict[str, Any]) -> Dict[str, Any]:
        """生成报告摘要"""
        try:
            summary = {
                'data_points': len(report_data),
                'has_errors': 'error' in report_data,
                'generated_at': datetime.utcnow().isoformat()
            }
            
            # 根据报告类型添加特定摘要
            if 'total_transactions' in report_data:
                summary['key_metrics'] = {
                    'transaction_count': report_data.get('total_transactions'),
                    'success_rate': report_data.get('success_rate')
                }
            elif 'total_violations' in report_data:
                summary['key_metrics'] = {
                    'violation_count': report_data.get('total_violations')
                }
            
            return summary
            
        except Exception as e:
            logger.error(f"Failed to generate report summary: {str(e)}")
            return {'error': str(e)}
    
    def get_audit_logs(self, user_id: str = None, action_type: str = None, 
                      start_date: datetime = None, end_date: datetime = None,
                      page: int = 1, page_size: int = 50) -> Dict[str, Any]:
        """获取审计日志"""
        try:
            query = {}
            
            if user_id:
                query['user_id'] = user_id
            if action_type:
                query['action_type'] = action_type
            if start_date or end_date:
                query['timestamp'] = {}
                if start_date:
                    query['timestamp']['$gte'] = start_date
                if end_date:
                    query['timestamp']['$lte'] = end_date
            
            skip = (page - 1) * page_size
            
            cursor = self.audit_logs_collection.find(query).sort('timestamp', -1).skip(skip).limit(page_size)
            logs = list(cursor)
            
            # 转换ObjectId和时间戳
            for log in logs:
                log['_id'] = str(log['_id'])
                log['timestamp'] = log['timestamp'].isoformat()
            
            total_count = self.audit_logs_collection.count_documents(query)
            
            return {
                'logs': logs,
                'total_count': total_count,
                'page': page,
                'page_size': page_size,
                'total_pages': (total_count + page_size - 1) // page_size
            }
            
        except Exception as e:
            logger.error(f"Failed to get audit logs: {str(e)}")
            raise AuditException(f"Failed to retrieve audit logs: {str(e)}")
    
    def get_regulatory_report(self, report_id: str) -> Dict[str, Any]:
        """获取监管报告"""
        try:
            report = self.regulatory_reports_collection.find_one({'report_id': report_id})
            
            if not report:
                raise AuditException(f"Report not found: {report_id}")
            
            # 转换时间戳
            report['_id'] = str(report['_id'])
            report['created_at'] = report['created_at'].isoformat()
            report['start_date'] = report['start_date'].isoformat()
            report['end_date'] = report['end_date'].isoformat()
            
            return report
            
        except Exception as e:
            if isinstance(e, AuditException):
                raise
            logger.error(f"Failed to get regulatory report: {str(e)}")
            raise AuditException(f"Failed to retrieve regulatory report: {str(e)}")
    
    def record_privacy_access(self, user_id: str, data_type: str, access_purpose: str, 
                            accessed_by: str) -> str:
        """记录隐私数据访问"""
        try:
            access_record = {
                'access_id': self._generate_access_id(),
                'user_id': user_id,
                'data_type': data_type,
                'access_purpose': access_purpose,
                'accessed_by': accessed_by,
                'access_timestamp': datetime.utcnow(),
                'ip_address': None,  # 应该从请求中获取
                'user_agent': None   # 应该从请求中获取
            }
            
            result = self.privacy_records_collection.insert_one(access_record)
            access_id = access_record['access_id']
            
            # 记录审计事件
            self.log_audit_event(
                user_id=accessed_by,
                action_type='privacy_data_access',
                details={
                    'target_user_id': user_id,
                    'data_type': data_type,
                    'access_purpose': access_purpose,
                    'access_id': access_id
                }
            )
            
            logger.info(f"Privacy access recorded: {access_id}")
            return access_id
            
        except Exception as e:
            logger.error(f"Failed to record privacy access: {str(e)}")
            raise AuditException(f"Failed to record privacy access: {str(e)}")
    
    def _generate_access_id(self) -> str:
        """生成访问ID"""
        timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S')
        hash_value = hashlib.md5(timestamp.encode()).hexdigest()[:8]
        return f"ACCESS_{timestamp}_{hash_value}"